from sklearn.datasets import fetch_california_housing
from sklearn.model_selection import validation_curve
from sklearn.ensemble import GradientBoostingRegressor as GBDT
import matplotlib.pyplot as plt


def main():
    # 加载数据
    california = fetch_california_housing()
    X = california.data[:150]
    y = california.target[:150]

    # 设置参数范围
    param_range = range(20, 150, 5)

    # 计算训练和验证分数
    train_score, val_score = validation_curve(
        GBDT(max_depth=3), X, y,
        param_name='n_estimators',
        param_range=param_range,
        cv=5,
        scoring='neg_mean_squared_error',  # 可以选择其他评分方式，这里使用负均方误差
        n_jobs=-1  # 使用所有CPU核心进行并行计算，加速过程
    )

    # 计算均值和标准差
    train_mean = train_score.mean(axis=-1)  # 注意：由于评分是负MSE，所以取反
    train_std = train_score.std(axis=-1)
    val_mean = val_score.mean(axis=-1)
    val_std = val_score.std(axis=-1)

    # 创建子图
    fig, axs = plt.subplots(1, 2, figsize=(12, 6))

    # 绘制训练和验证曲线
    axs[0].plot(param_range, train_mean, label='Training score')
    axs[0].fill_between(param_range, train_mean - train_std, train_mean + train_std, alpha=0.2)
    axs[0].set_title('Training score')
    axs[0].set_xlabel('Number of estimators')
    axs[0].set_ylabel('Score')

    axs[1].plot(param_range, val_mean, label='Validation score')
    axs[1].fill_between(param_range, val_mean - val_std, val_mean + val_std, alpha=0.2)
    axs[1].set_title('Validation score')
    axs[1].set_xlabel('Number of estimators')
    axs[1].set_ylabel('Score')

    # 显示图例
    for ax in axs:
        ax.legend()

    plt.show()


if __name__ == '__main__':
   main()